{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Yet Another OpenPose "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "desktop-sj3un3d\\bfgydx\n",
      "Python running from: E:\\Anaconda3\\envs\\YAOI\n",
      "Current working dir E:\\DeepLearning\\Yet-Another-Openpose-Implementation-master\\training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'gcloud' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n"
     ]
    }
   ],
   "source": [
    "#verify user\n",
    "!whoami\n",
    "#verify user account, if running on google cloud, otherwise ignore\n",
    "!gcloud auth list\n",
    "#which Environment/virtualenv running in\n",
    "import sys\n",
    "import os\n",
    "\n",
    "print(\"Python running from:\",sys.prefix)\n",
    "print(\"Current working dir\",os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "#start tensor board\n",
    "# must run \n",
    "#/usr/local/bin/tensorboard serve --logdir gs://dl_training_results/tensorboard --port 8889 --bind_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'gcloud' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n"
     ]
    }
   ],
   "source": [
    "#start TPU\n",
    "!gcloud compute tpus start node-1 --zone us-central1-c\n",
    "#stop TPU\n",
    "#!gcloud compute tpus stop node-1 --zone us-central1-c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (Temp/ipykernel_5512/1412333348.py, line 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  File \u001b[1;32m\"C:\\Users\\BFGYDX\\AppData\\Local\\Temp/ipykernel_5512/1412333348.py\"\u001b[1;36m, line \u001b[1;32m1\u001b[0m\n\u001b[1;33m    ---\u001b[0m\n\u001b[1;37m       ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "\n",
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"E:/DeepLearning/Yet-Another-Openpose-Implementation-master/configs/\")\n",
    "\n",
    "import default_config as cfg\n",
    "\n",
    "import local_storage_config as storage_cfg\n",
    "#import remote_storage_config as storage_cfg\n",
    "cfg.__dict__.update(storage_cfg.__dict__)\n",
    "\n",
    "import tpu_training.TPU_config as TPU_config\n",
    "cfg.__dict__.update(TPU_config.__dict__) #comment out to disable TPU training\n",
    "\n",
    "cfg.RUN_NAME=\"model11\" #for reference "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import now\n",
    "nowt=now()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version: 2.1.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "print(\"Tensorflow version:\",tf.version.VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%% \n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing results bucket connectivity\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'touch' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n",
      "'gsutil' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n",
      "'gsutil' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n",
      "'gsutil' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing dataset bucket connectivity\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'gsutil' 不是内部或外部命令，也不是可运行的程序\n",
      "或批处理文件。\n",
      "'gsutil' 不是内部或外部命令，也不是可运行的程序"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing TPU connectivity\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "或批处理文件。\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trying to connect to a TPU node\n",
      "\n",
      "!!!MAKE SURE THE TPU ADDRESS IS CORRECT!!\n",
      "1.ip must be correct\n",
      "2.tpu must be turned on\n",
      "3.version must be 'nightly-2.x'\n",
      "4.tpu must be reachable (check with gce networking/connectivity test)\n",
      "if not this will hang!\n",
      "\n",
      "Trying to connect to: grpc://10.0.3.2:8470\n"
     ]
    }
   ],
   "source": [
    "if cfg.TPU_MODE:\n",
    "    import tpu_training.init_TPU as init_TPU\n",
    "    \n",
    "    print(\"Testing results bucket connectivity\")\n",
    "    !touch /tmp/test\n",
    "    !gsutil cp /tmp/test {cfg.TENSORBOARD_PATH}/test\n",
    "    !gsutil rm {cfg.TENSORBOARD_PATH}/test\n",
    "    !gsutil cp /tmp/test {cfg.CHECKPOINTS_PATH}/test\n",
    "    !gsutil rm {cfg.CHECKPOINTS_PATH}/test\n",
    "    print(\"Testing dataset bucket connectivity\")\n",
    "    !gsutil ls gs://{cfg.GCS_TFRECORDS_BUCKETNAME} | head -4\n",
    "    print(\"Testing TPU connectivity\")\n",
    "    !nmap -Pn -p8470 {cfg.TPU_IP}\n",
    "    strategy,resolver=init_TPU.init_tpu(cfg.TPU_IP) #This must be run before any imports!!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# import datetime\n",
    "import numpy as np\n",
    "\n",
    "import dataset_functions\n",
    "import models.six_stage_linear_model as model\n",
    "import callbacks\n",
    "import dataset_builder\n",
    "import load_weights\n",
    "import loss_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# import importlib as il\n",
    "# il.reload() #useful to reload any changes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training settings\n",
    "TRAINING_EPOCHS = 30\n",
    "REAL_EPOCH_STEPS = int(cfg.DATASET_SIZE / cfg.BATCH_SIZE)\n",
    "SHORT_EPOCH_STEPS=50 #actual epocsh used in training, smaller than real epoch, but allows to track progress better, [in batches]\n",
    "SHORT_TRAINING_EPOCHS=int(TRAINING_EPOCHS*(REAL_EPOCH_STEPS/SHORT_EPOCH_STEPS))\n",
    "SHORT_VALIDATION_STEPS=5 #per short epoch\n",
    "\n",
    "EPOCH_RATIO=int(REAL_EPOCH_STEPS / SHORT_EPOCH_STEPS)\n",
    "# adam_learning_rate=0.001  #for reference\n",
    "BASE_LEARNING_RATE = 0.001\n",
    "LEARNING_RATE_SCHEDUELE = np.zeros(100000)  #used with short epochs\n",
    "LEARNING_RATE_SCHEDUELE[:3 * EPOCH_RATIO] = 1\n",
    "LEARNING_RATE_SCHEDUELE[3 * EPOCH_RATIO:20 * EPOCH_RATIO] = 1\n",
    "LEARNING_RATE_SCHEDUELE[20 * EPOCH_RATIO:40 * EPOCH_RATIO] = 1\n",
    "LEARNING_RATE_SCHEDUELE[40 * EPOCH_RATIO:100 * EPOCH_RATIO] = 0.5\n",
    "LEARNING_RATE_SCHEDUELE[100 * EPOCH_RATIO:] = 0.3\n",
    "LEARNING_RATE_SCHEDUELE *= BASE_LEARNING_RATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "model_ds=model.ModelDatasetComponent(cfg)\n",
    "\n",
    "tfrecord_files_train=dataset_builder.get_tfrecord_filenames(cfg.TRAIN_TFRECORDS,cfg)\n",
    "tfrecord_files_valid=dataset_builder.get_tfrecord_filenames(cfg.VALID_TFRECORDS,cfg)\n",
    "print(\"Found the following training TFrecords:\\n\",\"\\n\".join(tfrecord_files_train))\n",
    "print(\"Found the following validation TFrecords:\\n\",\"\\n\".join(tfrecord_files_valid))\n",
    "\n",
    "print(\"Building training dataset\")\n",
    "dst=dataset_builder.build_training_ds(tfrecord_files_train,model_ds.place_training_labels,cfg)\n",
    "print(\"Training dataset shape:\",dst)\n",
    "print(\"Building validation dataset\")\n",
    "dsv=dataset_builder.build_validation_ds(tfrecord_files_valid,model_ds.place_training_labels,cfg)\n",
    "print(\"Validation dataset shape:\",dsv)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#Test element\n",
    "# dst_iter=iter(dst)\n",
    "# sample_elem=next(dst_iter)\n",
    "# print(\"Dataset shape:\",dst) #this should match the model input, and output stages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#testing mask\n",
    "# sample_elem=next(dst_iter)\n",
    "# m=sample_elem[1][1][0,...,0]\n",
    "# import matplotlib.pyplot as plt\n",
    "# plt.imshow(m)\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Model\n",
    "### Load weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "if cfg.ASK_FOR_CHECKPOINTS:\n",
    "    checkpoint,starting_epoch=load_weights.checkpoints_prompt(cfg)\n",
    "else:\n",
    "    checkpoint=None\n",
    "    starting_epoch=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_maker=model.ModelMaker(cfg) #must be outside scope to keep the graph clean\n",
    "tf.keras.backend.clear_session() #to clean to backaend from the imported model\n",
    "def define():\n",
    "    train_model,test_model=model_maker.create_models()\n",
    "    \n",
    "#     if cfg.INCLUDE_MASK:\n",
    "#         losses=[loss_metrics.MaskedMeanAbsoluteError()\n",
    "#                 ,loss_metrics.MaskedMeanAbsoluteError()\n",
    "#                 ,loss_metrics.MaskedMeanAbsoluteError()\n",
    "#                 ,loss_metrics.MaskedMeanAbsoluteError()\n",
    "#                 ,loss_metrics.MaskedMeanSquaredError()\n",
    "#                 ,loss_metrics.MaskedMeanSquaredError()]                      \n",
    "#     else:\n",
    "#         raise NotImplementedError       \n",
    "    \n",
    "    #this must match the model output order\n",
    "    metrics=[\n",
    "              [loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]\n",
    "             ,[loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]\n",
    "             ,[loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]\n",
    "             ,[loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]\n",
    "             ,[loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]    \n",
    "             ,[loss_metrics.MeanAbsoluteRatio(),loss_metrics.AnalogRecall()]\n",
    "        ]\n",
    "    \n",
    "    train_model.compile(\n",
    "                    optimizer=tf.keras.optimizers.Adam(BASE_LEARNING_RATE)                   \n",
    "                    ,loss=loss_metrics.MaskedMeanSquaredError()\n",
    "                    ,loss_weights=[5,5,5,5,1,1]\n",
    "                    ,metrics=metrics                           \n",
    "                   )\n",
    "    return train_model,test_model\n",
    "\n",
    "if cfg.TPU_MODE:\n",
    "    print(\"TPU\")\n",
    "    with strategy.scope():\n",
    "        train_model,test_model=define()\n",
    "        if (checkpoint):\n",
    "            train_model.load_weights(checkpoint)\n",
    "else:\n",
    "    print(\"gpu\")\n",
    "    train_model,test_model=define()\n",
    "    if (checkpoint):\n",
    "        train_model.load_weights(checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "all_callbacks=[\n",
    "    callbacks.make_LRscheduler_callback(LEARNING_RATE_SCHEDUELE)\n",
    "    ,callbacks.PrintLR()\n",
    "    ,tf.keras.callbacks.TerminateOnNaN()\n",
    "]\n",
    "if cfg.SAVE_CHECKPOINTS:\n",
    "    checkpoint_callback,checkpoint_path=callbacks.make_checkpoint_callback(cfg,nowt,REAL_EPOCH_STEPS*cfg.BATCH_SIZE)\n",
    "    all_callbacks.append(checkpoint_callback)\n",
    "if cfg.SAVE_TENSORBOARD:\n",
    "    tensorboard_callback,tensorboard_path=callbacks.make_tensorboard_callback(cfg,nowt,EPOCH_RATIO)\n",
    "    all_callbacks+=[tensorboard_callback]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "#run to clean all tensorboard&checkpoints results\n",
    "#!gsutil -m rm -r {cfg.TENSORBOARD_PATH}/*\n",
    "#!gsutil -m rm -r {cfg.CHECKPOINTS_PATH}/*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Training\n",
    "Actually training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_history=train_model.fit(\n",
    "    dst\n",
    "    ,epochs=SHORT_TRAINING_EPOCHS\n",
    "    ,steps_per_epoch=SHORT_EPOCH_STEPS\n",
    "    ,validation_steps=SHORT_VALIDATION_STEPS\n",
    "    ,validation_data=dsv\n",
    "    ,callbacks=all_callbacks\n",
    "    ,initial_epoch=starting_epoch\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#run to clear this run results\n",
    "#!gsutil -m rm -r {tensorboard_path}/*\n",
    "#!gsutil -m rm -r {checkpoint_path}/*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "tmp_path='gs://dl_training_results/tmp/tensorflow/temp_weights'\n",
    "train_model.save_weights(tmp_path)\n",
    "\n",
    "local_model_maker=model.ModelMaker(cfg) #must be outside scope to keep the graph clean\n",
    "tf.keras.backend.clear_session() #to clean to backaend from the imported model\n",
    "\n",
    "cpu_train_model,cpu_test_model=local_model_maker.create_models()\n",
    "\n",
    "cpu_train_model.load_weights(tmp_path)\n",
    "cpu_test_model.load_weights(tmp_path)\n",
    "\n",
    "cpu_test_model.save(cfg.MODELS_PATH+\"/test_\"+cfg.RUN_NAME+nowt,include_optimizer=False) \n",
    "cpu_train_model.save(cfg.MODELS_PATH+\"/train_\"+cfg.RUN_NAME+nowt,include_optimizer=False) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### Shutdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#shut down TPU\n",
    "!gcloud compute tpus stop node-1 --zone us-central1-c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "#shut down this host\n",
    "!gcloud compute instances stop instance-1 --zone us-central1-c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "原始单元格格式",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
